"""
A simple example for Reinforcement Learning using table lookup Q-learning method.
An agent "o" is on the left of a 1 dimensional world, the treasure is on the rightmost location.
Run this program and to see how the agent will improve its strategy of finding the treasure.

View more on my tutorial page: https://morvanzhou.github.io/tutorials/
"""

import numpy as np
import pandas as pd
import time

np.random.seed(2)  # reproducible


N_STATES = 6  # the length of the 1 dimensional world
ACTIONS = ["left", "right"]  # available actions
EPSILON = 0.9  # greedy police
ALPHA = 0.1  # learning rate
GAMMA = 0.9  # discount factor
MAX_EPISODES = 13  # maximum episodes
FRESH_TIME = 0.3  # fresh time for one move


def build_q_table(n_states, actions):
    # 创建一个空的Q表（Q-table）
    table = pd.DataFrame(
        np.zeros(
            (n_states, len(actions))
        ),  # 使用全零矩阵初始化Q表，大小为(n_states, len(actions))
        columns=actions,  # 使用动作列表作为列名
    )
    # print(table)    # 可选：打印Q表，用于显示初始值
    return table  # 返回构建好的Q表


def choose_action(state, q_table):
    # 从Q表中选择一个动作
    state_actions = q_table.iloc[state, :]  # 获取当前状态下各个动作的Q值
    if (np.random.uniform() > EPSILON) or ((state_actions == 0).all()):
        # 如果随机数大于EPSILON，或者当前状态下所有动作的Q值都为0（即探索），则选择非贪婪动作
        action_name = np.random.choice(ACTIONS)  # 在所有动作中随机选择一个
    else:
        # 否则选择贪婪动作（选择Q值最大的动作）
        action_name = state_actions.idxmax()  # 使用idxmax函数找到Q值最大的动作
        # 注意：在较新版本的pandas中，argmax函数已被替换为idxmax
    return action_name  # 返回选择的动作名称


# 模拟Agent与环境的交互过程，根据Agent采取的动作以及当前状态，返回下一个状态和即时奖励。
def get_env_feedback(S, A):
    # Agent与环境进行交互的函数

    if A == "right":  # 如果动作是向右移动
        if S == N_STATES - 2:  # 如果当前状态是倒数第二个状态（即接近终点）
            S_ = "terminal"  # 下一个状态是终止状态
            R = 1  # 获得奖励1
        else:  # 如果当前状态不是倒数第二个状态
            S_ = S + 1  # 下一个状态向右移动一步
            R = 0  # 没有奖励
    else:  # 如果动作是向左移动
        R = 0  # 没有奖励
        if S == 0:  # 如果当前状态是起点位置
            S_ = S  # 下一个状态还是起点位置（到达墙壁无法再向左移动）
        else:  # 如果当前状态不是起点位置
            S_ = S - 1  # 下一个状态向左移动一步
    return S_, R  # 返回下一个状态和奖励值


# 根据Agent的当前状态，在环境中更新Agent的位置并打印出当前环境状态
def update_env(S, episode, step_counter):
    # 更新环境的函数
    env_list = ["-"] * (N_STATES - 1) + ["T"]  # '---------T' 是我们的环境
    if S == "terminal":  # 如果当前状态是终止状态
        interaction = "Episode %s: total_steps = %s" % (episode + 1, step_counter)
        print("\r{}".format(interaction), end="")  # 打印当前情节和总步数
        time.sleep(2)  # 等待2秒
        print("\r                                ", end="")  # 清空当前行
    else:  # 如果当前状态不是终止状态
        env_list[S] = "o"  # 在当前状态的位置上放置代表Agent的符号
        interaction = "".join(env_list)  # 将环境列表转换为字符串
        print("\r{}".format(interaction), end="")  # 打印当前环境状态
        time.sleep(FRESH_TIME)  # 等待一段时间，使得环境更新可见


# 强化学习的主要循环部分。它包括了强化学习的核心逻辑，使用Q-learning算法来更新Q表格。
def rl():
    # 强化学习的主要循环部分
    q_table = build_q_table(N_STATES, ACTIONS)  # 初始化Q表格
    for episode in range(MAX_EPISODES):  # 循环每一轮episode
        step_counter = 0  # 步数计数器初始化为0
        S = 0  # 初始状态设为0
        is_terminated = False  # 终止标志初始化为False
        update_env(S, episode, step_counter)  # 更新环境状态
        while not is_terminated:  # 当未终止时

            A = choose_action(S, q_table)  # 选择动作
            S_, R = get_env_feedback(S, A)  # 执行动作并获得下一个状态和奖励
            q_predict = q_table.loc[S, A]  # 用于预测的Q值
            if S_ != "terminal":  # 如果下一个状态不是终止状态
                q_target = (
                    R + GAMMA * q_table.iloc[S_, :].max()
                )  # 目标Q值（考虑未来奖励）
            else:  # 如果下一个状态是终止状态
                q_target = R  # 目标Q值为当前奖励
                is_terminated = True  # 标记为终止状态

            q_table.loc[S, A] += ALPHA * (q_target - q_predict)  # 更新Q表格
            S = S_  # 进入下一个状态

            update_env(S, episode, step_counter + 1)  # 更新环境状态
            step_counter += 1  # 步数加1
    return q_table  # 返回训练完成的Q表格


if __name__ == "__main__":
    q_table = rl()
    print("\r\nQ-table:\n")
    print(q_table)
